package random

import (
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/resolver"
	"log"
	"math"
	"sync"
	"sync/atomic"
)

type Builder struct {
}

func (b *Builder) Build(info base.PickerBuildInfo) balancer.Picker {
	connList := make([]*Conn, 0, len(info.ReadySCs))
	for subConn, sc := range info.ReadySCs {
		connList = append(connList, &Conn{
			conn:      subConn,
			addr:      sc.Address,
			activeCnt: math.MaxInt,
		})
	}

	bl := &Balancer{
		connList: connList,
	}

	return bl
}

type Balancer struct {
	connList []*Conn
	mu       sync.Mutex
}

func (b *Balancer) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(b.connList) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	// 最小活跃数
	b.mu.Lock()
	defer b.mu.Unlock()

	conn := b.connList[0]
	for _, c := range b.connList {
		if c.activeCnt <= conn.activeCnt {
			conn = c
		}
	}

	atomic.AddInt64(&conn.activeCnt, 1)

	return balancer.PickResult{
		SubConn: conn.conn,
		Done: func(info balancer.DoneInfo) {
			atomic.AddInt64(&conn.activeCnt, -1)
			if info.Err != nil {
				log.Println(info.Err)
				return
			}
		},
	}, nil
}

type Conn struct {
	conn      balancer.SubConn
	addr      resolver.Address
	activeCnt int64
}
